import os
import torch
import xml.etree.ElementTree as ET
import shutil
def get_config_from_xml(xml_path='config.xml'):
    tree = ET.parse(xml_path)
    root = tree.getroot()

    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    kwargs = {'pin_memory': True, 'num_workers': 4} if device == 'cuda' else {}

    image_size = tuple(map(int, root.find('image_size').text.split(', ')))

    output_task = os.path.join(root.find('output_path_root').text, os.path.basename(xml_path)[:-4])

    if not os.path.exists(output_task):
        os.makedirs(output_task)
        print(f"Task Directory '{output_task}' was created.")

        # 复制XML文件到新创建的目录
    destination_cfg_file_path = os.path.join(output_task, os.path.basename(xml_path))  # 目标文件的完整路径
    if not os.path.exists(destination_cfg_file_path):
        shutil.copy(xml_path, destination_cfg_file_path)
        print(f"File '{xml_path}' was copied to '{output_task}'.")
    model_name = root.find('model_name').text
    config = {
        'DEVICE': device,
        'KWARGS': kwargs,
        'IMAGE_PATH': os.path.join(root.find('data_path').text, 'images'),
        'DATASET_PATH': root.find('data_path').text,
        'MASK_PATH': os.path.join(root.find('data_path').text, 'annotations', 'trimaps'),
        'IMAGE_SIZE': image_size,
        'MODEL_NAME': root.find('model_name').text,
        'LOSS_FXN': root.find('loss_fxn').text,
        'N_CLASSES': int(root.find('n_classes').text),
        'BATCH_SIZE': int(root.find('batch_size').text),
        'RANDOM_SEED': int(root.find('random_seed').text),
        'SPLIT_RATE_TEST': float(root.find('split_rate_train_pred').text),
        'SPLIT_RATE_VALID': float(root.find('split_rate_train_valid').text),
        'IN_CHANNELS': int(root.find('in_channels').text),
        'OUT_CHANNELS': int(root.find('out_channels').text),
        'LEARNING_RATE': float(root.find('learning_rate').text),
        'EPOCH_NUM': int(root.find('epochs').text),
        'CP_OVERWRITE': int(root.find('overwrite_checkpoint').text),
        'OUTPUT_PATH': output_task,
        'SPLIT_FILE_IDX': os.path.join(output_task, 'split_file_idx.pickle'),
        'MODEL_PATH': os.path.join(output_task, f'{model_name}_model.pth'),
        'CHECKPOINT_PATH': os.path.join(output_task, f'{model_name}_checkpoint_model.pth'),
        'BEST_MODEL_PATH': os.path.join(output_task, f'{model_name}_best_model.pth'),
        'HISTORY_PATH': os.path.join(output_task, f'{model_name}_history.pickle'),
        'HISTORY_PLOT_PATH': os.path.join(output_task, f'{model_name}_train_loss.png'),
        'PRED_PLOT_PATH': os.path.join(output_task, f'{model_name}_pred')
    }

    return config

# Example of how to use the function
if "__main__" == __name__:
    config = get_config_from_xml('configurations/config_unet_bce_epoch50.xml')
    print(config)
